nanopyx.liquid._le_mandelbrot_benchmark_

  1import warnings
  2
  3import numpy as np
  4
  5try:
  6    from numba import njit, prange
  7
  8except ImportError:
  9    # raise a warning that numba is not installed
 10    # and that the njit functions will not be used
 11    # and that the pure python functions will be used instead
 12
 13    prange = range
 14
 15    def njit(*args, **kwargs):
 16        def wrapper(func):
 17            warnings.warn(
 18                f"Numba is not installed. Using pure python for {func.__name__}"
 19            )
 20            return func
 21
 22        return wrapper
 23
 24
 25MAX_ITERATIONS = 1000
 26DIVERGENCE = 10
 27
 28
 29def _py_mandelbrot(row: float, col: float) -> int:
 30    zrow = 0
 31    zcol = 0
 32    iterations = 0
 33    while zrow * zrow + zcol * zcol <= DIVERGENCE and iterations < MAX_ITERATIONS:
 34        zrow_new = zrow * zrow - zcol * zcol + row
 35        zcol_new = 2 * zrow * zcol + col
 36        zrow = zrow_new
 37        zcol = zcol_new
 38        iterations += 1
 39
 40    return iterations
 41
 42
 43@njit(cache=True)
 44def _njit_mandelbrot(row: float, col: float) -> int:
 45    zrow = 0
 46    zcol = 0
 47    iterations = 0
 48    while zrow * zrow + zcol * zcol <= DIVERGENCE and iterations < MAX_ITERATIONS:
 49        zrow_new = zrow * zrow - zcol * zcol + row
 50        zcol_new = 2 * zrow * zcol + col
 51        zrow = zrow_new
 52        zcol = zcol_new
 53        iterations += 1
 54
 55    return iterations
 56
 57
 58def mandelbrot(
 59    image: np.ndarray,
 60    r_start: float,
 61    r_end: float,
 62    c_start: float,
 63    c_end: float,
 64) -> np.ndarray:
 65    """
 66    Mandelbrot set generator.
 67    :param image: numpy array to store the result
 68    :param r_start: start of the real axis
 69    :param r_end: end of the real axis
 70    :param c_start: start of the complex axis
 71    :param c_end: end of the complex axis
 72    :return: numpy array with the result
 73    """
 74    rows, cols = image.shape
 75    for row in range(rows):
 76        for col in range(cols):
 77            image[row, col] = _py_mandelbrot(
 78                r_start + (r_end - r_start) * row / rows,
 79                c_start + (c_end - c_start) * col / cols,
 80            )
 81    return image
 82
 83
 84@njit(cache=True, parallel=True)
 85def njit_mandelbrot(
 86    image: np.ndarray,
 87    r_start: float,
 88    r_end: float,
 89    c_start: float,
 90    c_end: float,
 91) -> np.ndarray:
 92    """
 93    Mandelbrot set generator.
 94    :param image: numpy array to store the result
 95    :param r_start: start of the real axis
 96    :param r_end: end of the real axis
 97    :param c_start: start of the complex axis
 98    :param c_end: end of the complex axis
 99    :return: numpy array with the result
100    """
101    rows, cols = image.shape
102    for row in prange(rows):
103        for col in range(cols):
104            image[row, col] = _njit_mandelbrot(
105                r_start + (r_end - r_start) * row / rows,
106                c_start + (c_end - c_start) * col / cols,
107            )
108    return image